package com.eudon.ai.agent.etl.extension;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * @author eudon
 * @description 查询重写器
 * @date 2025/8/16
 * @email syd19970616@gmail.com
 * @address <a href="https://gitee.com/eudon">Gitee</a>
 */
@Component
@Slf4j
public class MyQueryRewriter {
    private final RewriteQueryTransformer queryTransformer;

    public MyQueryRewriter(ChatModel qwenTurbo) {
        this.queryTransformer =RewriteQueryTransformer.builder()
                .chatClientBuilder(ChatClient.builder(qwenTurbo))
                .build();
    }

    public Query rewrite(Query query) {
        return queryTransformer.transform(query);
    }
}

